package org.qh.ai.service.request.impl;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.annotation.PreDestroy;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import okio.BufferedSource;
import org.qh.ai.domain.ChatModel;
import org.qh.ai.handler.StreamingChatResponseHandler;
import org.qh.ai.service.request.ChatService;
import org.qh.common.ai.entity.chat.Message;
import org.qh.common.ai.entity.request.ChatRequest;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

@Slf4j
@Service
@RequiredArgsConstructor
public class OpenAiChatServiceImpl implements ChatService {

    private final ObjectMapper objectMapper;
    private static final String API_PATH = "/v1/chat/completions";
    private static final String AUTH_HEADER = "Authorization";
    private static final String CONTENT_TYPE = "Content-Type";
    private static final String JSON_MEDIA_TYPE = "application/json";
    private static final String DATA_PREFIX = "data: ";
    private static final String DONE_MARKER = "[DONE]";
    private static final String SSE_CONTENT_TYPE = "text/event-stream";

    private final ExecutorService executorService = Executors.newCachedThreadPool();

    private final OkHttpClient client = new OkHttpClient.Builder()
        .connectTimeout(30, TimeUnit.SECONDS)
        .readTimeout(60, TimeUnit.SECONDS)
        .writeTimeout(30, TimeUnit.SECONDS)
        .build();

    private Request buildRequest(ChatModel chatModel, ChatRequest request) throws IOException {
        if (chatModel == null || chatModel.getApiHost() == null || chatModel.getApiKey() == null) {
            throw new IllegalArgumentException("ChatModel or its configuration cannot be null");
        }

        // 确保设置为流式请求
        request.setStream(request.getStream());

        String requestBody = objectMapper.writeValueAsString(request);
        HttpUrl url = HttpUrl.parse(chatModel.getApiHost() + API_PATH);

        if (url == null) {
            throw new IllegalArgumentException("Invalid API host URL");
        }

        return new Request.Builder()
            .url(url)
            .addHeader(AUTH_HEADER, "Bearer " + chatModel.getApiKey())
            .addHeader(CONTENT_TYPE, JSON_MEDIA_TYPE)
            .post(RequestBody.create(requestBody, MediaType.get(JSON_MEDIA_TYPE)))
            .build();
    }

    private void handleStreamingResponse(Response response, StreamingChatResponseHandler handler) {
        try (ResponseBody body = response.body()) {
            if (body == null) {
                handler.onError(new IOException("Empty response body"));
                return;
            }
            BufferedSource source = body.source();
            StringBuilder completeResponse = new StringBuilder();
            try {
                while (!source.exhausted()) {
                    String line = source.readUtf8Line();
                    if (line == null) break;

                    line = line.trim();
                    if (line.isEmpty()) continue;

                    if (line.startsWith(DATA_PREFIX)) {
                        String data = line.substring(DATA_PREFIX.length());
                        // 解析并简化数据
                        String simplifyData = parseAndSimplify(data);
                        if (data.equals(DONE_MARKER)) {
                            // 响应结束
                            handler.onCompleteResponse(completeResponse.toString());
                            break;
                        }
                        if (Objects.nonNull(simplifyData)) {
                            // 收集所有响应数据
                            completeResponse.append(simplifyData);
                            // 发送部分响应数据
                            handler.onPartialResponse(simplifyData);
                        }
                    }
                }
            } catch (IOException e) {
                handler.onError(e);
                log.error("Error reading stream", e);
            }
        }
    }

    @Override
    public void chat(ChatModel chatModel, List<Message> messages, StreamingChatResponseHandler handler) {
        executorService.submit(() -> {
            try {
                ChatRequest request = new ChatRequest(messages, chatModel.getModelName());
                Request httpRequest = buildRequest(chatModel, request);

                try (Response response = client.newCall(httpRequest).execute()) {
                    if (!response.isSuccessful()) {
                        String errorBody = response.body() != null ? response.body().string() : "";
                        handler.onError(new IOException(
                            "Request failed with code: " + response.code() + ", body: " + errorBody));
                        return;
                    }

                    // 验证Content-Type
                    String contentType = response.header(CONTENT_TYPE);
                    if (contentType == null || !contentType.contains(SSE_CONTENT_TYPE)) {
                        handler.onError(new IOException("Expected text/event-stream but got: " + contentType));
                        return;
                    }

                    handleStreamingResponse(response, handler);
                }
            } catch (Exception e) {
                handler.onError(e);
                log.error("Chat request failed", e);
            }
        });
    }

    @Override
    public String chat(ChatModel chatModel, List<Message> messages) {
        try {
            ChatRequest request = new ChatRequest(messages, chatModel.getModelName());
            request.setStream(false); // 明确设置为非流式

            Request httpRequest = buildRequest(chatModel, request);

            try (Response response = client.newCall(httpRequest).execute()) {
                if (!response.isSuccessful()) {
                    log.error("Request failed with status: {}", response.code());
                    return "";
                }

                ResponseBody body = response.body();
                return body != null ? body.string() : "";
            }
        } catch (Exception e) {
            log.error("Chat failed", e);
            return "";
        }
    }

    private static final ObjectMapper mapper = new ObjectMapper();

    private static boolean isFirstThought = true;  // 标记是否是第一个思维内容
    private static final String THINK_START = "<think>";
    private static final String THINK_END = "</think>";

    /**
     * 从原始响应中提取content/reasoning_content内容
     * @param originalJson 原始JSON字符串
     * @return 提取的内容字符串，优先返回content，其次reasoning_content，解析失败返回null
     */
    public static String parseAndSimplify(String originalJson) {
        if (originalJson == null || originalJson.isEmpty() || DONE_MARKER.equals(originalJson)) {
            return null;
        }

        try {
            JsonNode jsonNode = mapper.readTree(originalJson);
            JsonNode choices = jsonNode.path("choices");

            if (choices.isArray() && !choices.isEmpty()) {
                JsonNode delta = choices.get(0).path("delta");

                // 1. 处理思维链内容 (reasoning_content)
                if (delta.has("reasoning_content")) {
                    String reasoningContent = delta.get("reasoning_content").asText();
                    if (!reasoningContent.isEmpty() && !"null".equals(reasoningContent)) {
                        String result = reasoningContent;
                        if (isFirstThought) {
                            result = THINK_START + result;
                            isFirstThought = false;
                        }
                        return result;
                    }
                }

                // 2. 处理最终回答 (content) - 遇到content时结束思维链
                if (delta.has("content")) {
                    String content = delta.get("content").asText();
                    if (!content.isEmpty() && !"null".equals(content)) {
                        if (!isFirstThought) {  // 如果有未关闭的思维链
                            content = THINK_END + content;
                            isFirstThought = true;  // 重置状态
                        }
                        return content;
                    }
                }
            }
        } catch (Exception e) {
            log.warn("Failed to parse response chunk: {}", originalJson, e);
        }

        return null;
    }

    @PreDestroy
    public void shutdown() {
        executorService.shutdown();
        try {
            if (!executorService.awaitTermination(5, TimeUnit.SECONDS)) {
                executorService.shutdownNow();
            }
        } catch (InterruptedException e) {
            executorService.shutdownNow();
            Thread.currentThread().interrupt();
        }
    }
}
